#include <stdio.h>
#include <string.h>
#include <openssl/evp.h>
#include <openssl/pem.h>
#include <openssl/ec.h>
//#include <openssl/sm2.h> // SM2 特定头文件

// (复用 print_hex 函数)
void print_hex(const char* label, const unsigned char* data, int len) {
    printf("%s (%d bytes):\n", label, len);
    for (int i = 0; i < len; ++i) {
        printf("%02x", data[i]);
    }
    printf("\n");
}

// SM2 密钥对
EVP_PKEY *create_sm2_key() {
    EVP_PKEY *pkey = NULL;
    EVP_PKEY_CTX *ctx = EVP_PKEY_CTX_new_id(EVP_PKEY_SM2, NULL);
    if (ctx == NULL) return NULL;

    if (EVP_PKEY_keygen_init(ctx) <= 0) {
        EVP_PKEY_CTX_free(ctx);
        return NULL;
    }
    
    if (EVP_PKEY_keygen(ctx, &pkey) <= 0) {
        EVP_PKEY_CTX_free(ctx);
        return NULL;
    }
    
    EVP_PKEY_CTX_free(ctx);
    return pkey;
}

int main() {
    EVP_PKEY *pkey = NULL;
    EVP_PKEY_CTX *ctx;
    unsigned char *ciphertext, *decryptedtext;
    size_t ciphertext_len, decryptedtext_len;

    const unsigned char msg[] = "Test message for SM2";
    
    // 1. 生成 SM2 密钥对
    pkey = create_sm2_key();
    if (pkey == NULL) {
        fprintf(stderr, "Failed to create SM2 key\n");
        return 1;
    }
    printf("SM2 Key generation successful.\n");

    // 2. SM2 加密
    ctx = EVP_PKEY_CTX_new(pkey, NULL);
    if (EVP_PKEY_encrypt_init(ctx) <= 0) {
        fprintf(stderr, "EVP_PKEY_encrypt_init failed\n");
        return 1;
    }

    // 确定密文所需缓冲区大小
    if (EVP_PKEY_encrypt(ctx, NULL, &ciphertext_len, msg, strlen((char*)msg)) <= 0) {
        fprintf(stderr, "EVP_PKEY_encrypt (size) failed\n");
        return 1;
    }
    
    ciphertext = OPENSSL_malloc(ciphertext_len);
    if (ciphertext == NULL) {
        fprintf(stderr, "OPENSSL_malloc failed\n");
        return 1;
    }

    // 执行加密
    if (EVP_PKEY_encrypt(ctx, ciphertext, &ciphertext_len, msg, strlen((char*)msg)) <= 0) {
        fprintf(stderr, "EVP_PKEY_encrypt failed\n");
        return 1;
    }
    
    print_hex("SM2 Ciphertext", ciphertext, ciphertext_len);
    EVP_PKEY_CTX_free(ctx);

    // 3. SM2 解密
    ctx = EVP_PKEY_CTX_new(pkey, NULL);
    if (EVP_PKEY_decrypt_init(ctx) <= 0) {
        fprintf(stderr, "EVP_PKEY_decrypt_init failed\n");
        return 1;
    }

    // 确定明文所需缓冲区大小
    if (EVP_PKEY_decrypt(ctx, NULL, &decryptedtext_len, ciphertext, ciphertext_len) <= 0) {
        fprintf(stderr, "EVP_PKEY_decrypt (size) failed\n");
        return 1;
    }
    
    decryptedtext = OPENSSL_malloc(decryptedtext_len);
    if (decryptedtext == NULL) {
        fprintf(stderr, "OPENSSL_malloc failed\n");
        return 1;
    }

    // 执行解密
    if (EVP_PKEY_decrypt(ctx, decryptedtext, &decryptedtext_len, ciphertext, ciphertext_len) <= 0) {
        fprintf(stderr, "EVP_PKEY_decrypt failed\n");
        return 1;
    }

    decryptedtext[decryptedtext_len] = '\0';
    printf("SM2 Decrypted text: %s\n", decryptedtext);
    if (strcmp((char*)msg, (char*)decryptedtext) == 0) {
        printf("SM2 Enc/Dec Success!\n");
    }

    OPENSSL_free(ciphertext);
    OPENSSL_free(decryptedtext);
    EVP_PKEY_CTX_free(ctx);

    // ---------------------------------
    // 4. SM2 签名
    EVP_MD_CTX *md_ctx;
    unsigned char *sig;
    size_t sig_len;
    
    // SM2 签名需要指定 ID，这里使用默认的 "1234567812345678"
    // OpenSSL 3.x 默认使用 EVP_sm3() 进行哈希
    
    md_ctx = EVP_MD_CTX_new();
    if (md_ctx == NULL) return 1;

    // 使用 EVP_DigestSign* API
    if (EVP_DigestSignInit(md_ctx, &ctx, EVP_sm3(), NULL, pkey) <= 0) {
        fprintf(stderr, "EVP_DigestSignInit failed\n");
        return 1;
    }
    
    // 设置 SM2 签名所需的 ID
    if (EVP_PKEY_CTX_set1_id(ctx, (unsigned char *)"1234567812345678", 16) <= 0) {
        fprintf(stderr, "EVP_PKEY_CTX_set1_id failed\n");
        return 1;
    }

    if (EVP_DigestSignUpdate(md_ctx, msg, strlen((char*)msg)) <= 0) {
        fprintf(stderr, "EVP_DigestSignUpdate failed\n");
        return 1;
    }

    // 获取签名大小
    if (EVP_DigestSignFinal(md_ctx, NULL, &sig_len) <= 0) {
        fprintf(stderr, "EVP_DigestSignFinal (size) failed\n");
        return 1;
    }
    
    sig = OPENSSL_malloc(sig_len);
    if (sig == NULL) return 1;

    // 执行签名
    if (EVP_DigestSignFinal(md_ctx, sig, &sig_len) <= 0) {
        fprintf(stderr, "EVP_DigestSignFinal failed\n");
        return 1;
    }

    print_hex("SM2 Signature", sig, sig_len);

    // 5. SM2 验签
    // 重置 MD CTX
    if (EVP_DigestVerifyInit(md_ctx, &ctx, EVP_sm3(), NULL, pkey) <= 0) {
        fprintf(stderr, "EVP_DigestVerifyInit failed\n");
        return 1;
    }
    
    // 验签也需要同样的 ID
    if (EVP_PKEY_CTX_set1_id(ctx, (unsigned char *)"1234567812345678", 16) <= 0) {
        fprintf(stderr, "EVP_PKEY_CTX_set1_id (verify) failed\n");
        return 1;
    }

    if (EVP_DigestVerifyUpdate(md_ctx, msg, strlen((char*)msg)) <= 0) {
        fprintf(stderr, "EVP_DigestVerifyUpdate failed\n");
        return 1;
    }

    if (EVP_DigestVerifyFinal(md_ctx, sig, sig_len) == 1) {
        printf("SM2 Signature Verification Success!\n");
    } else {
        printf("SM2 Signature Verification Failed!\n");
    }

    EVP_PKEY_free(pkey);
    EVP_MD_CTX_free(md_ctx);
    OPENSSL_free(sig);

    return 0;
}
